{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\n# Installs the latest dev build of TVM from PyPI. If you wish to build\n# from source, see https://tvm.apache.org/docs/install/from_source.html\npip install apache-tvm --pre"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# 6. Model Tuning with microTVM\n**Authors**:\n[Andrew Reusch](https://github.com/areusch),\n[Mehrdad Hessar](https://github.com/mehrdadh)\n\nThis tutorial explains how to autotune a model using the C runtime.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Install microTVM Python dependencies\n\nTVM does not include a package for Python serial communication, so\nwe must install one before using microTVM. We will also need TFLite\nto load models.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\npip install pyserial==3.5 tflite==2.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# You can skip the following section (installing Zephyr) if the following flag is False.\n# Installing Zephyr takes ~20 min.\nimport os\n\nuse_physical_hw = bool(os.getenv(\"TVM_MICRO_USE_HW\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Install Zephyr\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\n# Install west and ninja\npython3 -m pip install west\napt-get install -y ninja-build\n\n# Install ZephyrProject\nZEPHYR_PROJECT_PATH=\"/content/zephyrproject\"\nexport ZEPHYR_BASE=${ZEPHYR_PROJECT_PATH}/zephyr\nwest init ${ZEPHYR_PROJECT_PATH}\ncd ${ZEPHYR_BASE}\ngit checkout v3.2-branch\ncd ..\nwest update\nwest zephyr-export\nchmod -R o+w ${ZEPHYR_PROJECT_PATH}\n\n# Install Zephyr SDK\ncd /content\nZEPHYR_SDK_VERSION=\"0.15.2\"\nwget \"https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v${ZEPHYR_SDK_VERSION}/zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\ntar xvf \"zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\nmv \"zephyr-sdk-${ZEPHYR_SDK_VERSION}\" zephyr-sdk\nrm \"zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\n\n# Install python dependencies\npython3 -m pip install -r \"${ZEPHYR_BASE}/scripts/requirements.txt\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import Python dependencies\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import json\nimport numpy as np\nimport pathlib\n\nimport tvm\nfrom tvm.relay.backend import Runtime\nimport tvm.micro.testing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Defining the model\n\n To begin with, define a model in Relay to be executed on-device. Then create an IRModule from relay model and\n fill parameters with random numbers.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_shape = (1, 3, 10, 10)\nweight_shape = (6, 3, 5, 5)\n\ndata = tvm.relay.var(\"data\", tvm.relay.TensorType(data_shape, \"float32\"))\nweight = tvm.relay.var(\"weight\", tvm.relay.TensorType(weight_shape, \"float32\"))\n\ny = tvm.relay.nn.conv2d(\n    data,\n    weight,\n    padding=(2, 2),\n    kernel_size=(5, 5),\n    kernel_layout=\"OIHW\",\n    out_dtype=\"float32\",\n)\nf = tvm.relay.Function([data, weight], y)\n\nrelay_mod = tvm.IRModule.from_expr(f)\nrelay_mod = tvm.relay.transform.InferType()(relay_mod)\n\nweight_sample = np.random.rand(\n    weight_shape[0], weight_shape[1], weight_shape[2], weight_shape[3]\n).astype(\"float32\")\nparams = {\"weight\": weight_sample}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Defining the target\n Now we define the TVM target that describes the execution environment. This looks very similar\n to target definitions from other microTVM tutorials. Alongside this we pick the C Runtime to code\n generate our model against.\n\n When running on physical hardware, choose a target and a board that\n describe the hardware. There are multiple hardware targets that could be selected from\n PLATFORM list in this tutorial. You can chose the platform by passing --platform argument when running\n this tutorial.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "RUNTIME = Runtime(\"crt\", {\"system-lib\": True})\nTARGET = tvm.micro.testing.get_target(\"crt\")\n\n# Compiling for physical hardware\n# --------------------------------------------------------------------------\n#  When running on physical hardware, choose a TARGET and a BOARD that describe the hardware. The\n#  STM32L4R5ZI Nucleo target and board is chosen in the example below.\nif use_physical_hw:\n    BOARD = os.getenv(\"TVM_MICRO_BOARD\", default=\"nucleo_l4r5zi\")\n    SERIAL = os.getenv(\"TVM_MICRO_SERIAL\", default=None)\n    TARGET = tvm.micro.testing.get_target(\"zephyr\", BOARD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Extracting tuning tasks\n Not all operators in the Relay program printed above can be tuned. Some are so trivial that only\n a single implementation is defined; others don't make sense as tuning tasks. Using\n `extract_from_program`, you can produce a list of tunable tasks.\n\n Because task extraction involves running the compiler, we first configure the compiler's\n transformation passes; we'll apply the same configuration later on during autotuning.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pass_context = tvm.transform.PassContext(opt_level=3, config={\"tir.disable_vectorize\": True})\nwith pass_context:\n    tasks = tvm.autotvm.task.extract_from_program(relay_mod[\"main\"], {}, TARGET)\nassert len(tasks) > 0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Configuring microTVM\n Before autotuning, we need to define a module loader and then pass that to\n a `tvm.autotvm.LocalBuilder`. Then we create a `tvm.autotvm.LocalRunner` and use\n both builder and runner to generates multiple measurements for auto tunner.\n\n In this tutorial, we have the option to use x86 host as an example or use different targets\n from Zephyr RTOS. If you choose pass `--platform=host` to this tutorial it will uses x86. You can\n choose other options by choosing from `PLATFORM` list.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "module_loader = tvm.micro.AutoTvmModuleLoader(\n    template_project_dir=pathlib.Path(tvm.micro.get_microtvm_template_projects(\"crt\")),\n    project_options={\"verbose\": False},\n)\nbuilder = tvm.autotvm.LocalBuilder(\n    n_parallel=1,\n    build_kwargs={\"build_option\": {\"tir.disable_vectorize\": True}},\n    do_fork=True,\n    build_func=tvm.micro.autotvm_build_func,\n    runtime=RUNTIME,\n)\nrunner = tvm.autotvm.LocalRunner(number=1, repeat=1, timeout=100, module_loader=module_loader)\n\nmeasure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)\n\n# Compiling for physical hardware\nif use_physical_hw:\n    module_loader = tvm.micro.AutoTvmModuleLoader(\n        template_project_dir=pathlib.Path(tvm.micro.get_microtvm_template_projects(\"zephyr\")),\n        project_options={\n            \"board\": BOARD,\n            \"verbose\": False,\n            \"project_type\": \"host_driven\",\n            \"serial_number\": SERIAL,\n        },\n    )\n    builder = tvm.autotvm.LocalBuilder(\n        n_parallel=1,\n        build_kwargs={\"build_option\": {\"tir.disable_vectorize\": True}},\n        do_fork=False,\n        build_func=tvm.micro.autotvm_build_func,\n        runtime=RUNTIME,\n    )\n    runner = tvm.autotvm.LocalRunner(number=1, repeat=1, timeout=100, module_loader=module_loader)\n\n    measure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Run Autotuning\n Now we can run autotuning separately on each extracted task on microTVM device.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "autotune_log_file = pathlib.Path(\"microtvm_autotune.log.txt\")\nif os.path.exists(autotune_log_file):\n    os.remove(autotune_log_file)\n\nnum_trials = 10\nfor task in tasks:\n    tuner = tvm.autotvm.tuner.GATuner(task)\n    tuner.tune(\n        n_trial=num_trials,\n        measure_option=measure_option,\n        callbacks=[\n            tvm.autotvm.callback.log_to_file(str(autotune_log_file)),\n            tvm.autotvm.callback.progress_bar(num_trials, si_prefix=\"M\"),\n        ],\n        si_prefix=\"M\",\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Timing the untuned program\n For comparison, let's compile and run the graph without imposing any autotuning schedules. TVM\n will select a randomly-tuned implementation for each operator, which should not perform as well as\n the tuned operator.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with pass_context:\n    lowered = tvm.relay.build(relay_mod, target=TARGET, runtime=RUNTIME, params=params)\n\ntemp_dir = tvm.contrib.utils.tempdir()\nproject = tvm.micro.generate_project(\n    str(tvm.micro.get_microtvm_template_projects(\"crt\")),\n    lowered,\n    temp_dir / \"project\",\n    {\"verbose\": False},\n)\n\n# Compiling for physical hardware\nif use_physical_hw:\n    temp_dir = tvm.contrib.utils.tempdir()\n    project = tvm.micro.generate_project(\n        str(tvm.micro.get_microtvm_template_projects(\"zephyr\")),\n        lowered,\n        temp_dir / \"project\",\n        {\n            \"board\": BOARD,\n            \"verbose\": False,\n            \"project_type\": \"host_driven\",\n            \"serial_number\": SERIAL,\n            \"config_main_stack_size\": 4096,\n        },\n    )\n\nproject.build()\nproject.flash()\nwith tvm.micro.Session(project.transport()) as session:\n    debug_module = tvm.micro.create_local_debug_executor(\n        lowered.get_graph_json(), session.get_system_lib(), session.device\n    )\n    debug_module.set_input(**lowered.get_params())\n    print(\"########## Build without Autotuning ##########\")\n    debug_module.run()\n    del debug_module"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Timing the tuned program\n Once autotuning completes, you can time execution of the entire program using the Debug Runtime:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with tvm.autotvm.apply_history_best(str(autotune_log_file)):\n    with pass_context:\n        lowered_tuned = tvm.relay.build(relay_mod, target=TARGET, runtime=RUNTIME, params=params)\n\ntemp_dir = tvm.contrib.utils.tempdir()\nproject = tvm.micro.generate_project(\n    str(tvm.micro.get_microtvm_template_projects(\"crt\")),\n    lowered_tuned,\n    temp_dir / \"project\",\n    {\"verbose\": False},\n)\n\n# Compiling for physical hardware\nif use_physical_hw:\n    temp_dir = tvm.contrib.utils.tempdir()\n    project = tvm.micro.generate_project(\n        str(tvm.micro.get_microtvm_template_projects(\"zephyr\")),\n        lowered_tuned,\n        temp_dir / \"project\",\n        {\n            \"board\": BOARD,\n            \"verbose\": False,\n            \"project_type\": \"host_driven\",\n            \"serial_number\": SERIAL,\n            \"config_main_stack_size\": 4096,\n        },\n    )\n\nproject.build()\nproject.flash()\nwith tvm.micro.Session(project.transport()) as session:\n    debug_module = tvm.micro.create_local_debug_executor(\n        lowered_tuned.get_graph_json(), session.get_system_lib(), session.device\n    )\n    debug_module.set_input(**lowered_tuned.get_params())\n    print(\"########## Build with Autotuning ##########\")\n    debug_module.run()\n    del debug_module"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}